from openai import OpenAI
import os

# 加载环境变量
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) # 读取本地 .env 文件，里面定义了 OPENAI_API_KEY

client = OpenAI()

def get_completion(prompt, model="gpt-4o-mini"):
    ''' 获取 completions '''
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,
    )
    # print(response)
    return response.choices[0].message.content

def build_prompt(propmt_template, **kwargs):
    ''' 将 Prompt 模板与参数拼接成完整的 Prompt '''
    inputs = {}
    for k, v in kwargs.items():
        if isinstance(v, list) and all(isinstance(elem, str) for elem in v):
            val = '\n\n'.join(v)
        else:
            val = v
        inputs[k] = val
    return propmt_template.format(**inputs)

prompt_template = """
你是一个问答机器人。
你的任务是根据下述给定的已知信息回答用户问题。

已知信息:
{context}

用户问：
{query}

如果已知信息不包含用户问题的答案，或者已知信息不足以回答用户的问题，请直接回复"我无法回答您的问题"。
请不要输出已知信息中不包含的信息或答案。
请用中文回答用户问题。
"""

if __name__ == "__main__":
    user_query = "how many parameters does llama 2 have?"
    from rag02 import search
    # 1.检索 依赖 Elasticsearch 数据库
    search_results = search(user_query, 2)

    # 2. 构建prompt
    prompt = build_prompt(prompt_template, context=search_results, query=user_query)
    print("======Prompt======")
    print(prompt)

    # 3. 调用LLM
    response = get_completion(prompt)

    print("====回复====")
    print(response)

